from langchain_deepseek import ChatDeepSeek 
from langchain.chat_models import init_chat_model
from langchain_core.prompts import PromptTemplate

from dotenv import load_dotenv

load_dotenv()

def generate_pet_name(animal_type, pet_color):
    model = ChatDeepSeek(model="deepseek-chat", temperature = 1)
    prompt = PromptTemplate(
        template = "我有一只宠物{animal_type}，它的颜色是{pet_color}色的，我想给它起个酷一点的名字，请推荐5个名字",
        input_variables=["animal_type", "pet_color"]
    )

    chain = prompt | model
    response = chain.invoke({"animal_type": animal_type, "pet_color": pet_color})

    return response

if __name__ == '__main__':
    print(generate_pet_name("猫", "黑"))
